package com.fiehn.util.dbunit;

import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.XmlDataSet;

import java.io.FileOutputStream;
import java.sql.Connection;
import java.util.HashSet;
import java.util.Set;


/**
 * Created by IntelliJ IDEA.
 * User: pradeep
 * Date: Mar 3, 2010
 * Time: 1:50:12 PM
 * Reference : http://www.dbunit.org/faq.html#extract
 */

/**
 * This class is used to export one or many tables from a database to a flat XML dataset file
 */
public class DatabaseExport {

    public static void main(String[] args) throws Exception {
        // database connection
        Connection jdbcConnection = DBConnection.getConnection();
        exportFullDataBase(jdbcConnection, "compound-repository-test");

        Set<String> set = new HashSet<String>();

        set.add("compound");
        set.add("smiles");
        set.add("iupac");
        set.add("inchi_hash_key");
        set.add("synonym");
        set.add("dblinks");


        exportPartialDataSet(jdbcConnection, set, "compound");
    }

    /**
     * Export the full database in XML file
     *
     * @param jdbcConnection
     * @param fileName
     * @throws Exception
     */
    public static void exportFullDataBase(Connection jdbcConnection, String fileName) throws Exception {
        IDatabaseConnection connection = new DatabaseConnection(jdbcConnection);

        // full database export
        IDataSet fullDataSet = connection.createDataSet();
        XmlDataSet.write(fullDataSet, new FileOutputStream(fileName + ".xml"));
    }

    /**
     * Export the data for a perticuler table
     *
     * @param jdbcConnection
     * @param tableName
     * @param fileName
     * @throws Exception
     */
    public static void exportPartialDataSet(Connection jdbcConnection, Set<String> tableName, String fileName) throws Exception {
        IDatabaseConnection connection = new DatabaseConnection(jdbcConnection);

        // partial database export
        QueryDataSet partialDataSet = new QueryDataSet(connection);

        for (String s : tableName) {
            partialDataSet.addTable(s);
        }

        XmlDataSet.write(partialDataSet, new FileOutputStream(fileName + ".xml"));
    }


}
